import axios from 'axios';

export class BGEEmbeddings {
  client: axios.AxiosInstance;
  apiKey: string;
  apiUrl: string;
  modelName: string;

  constructor(apiKey: string, apiUrl: string, modelName: string) {
    this.apiKey = apiKey;
    this.apiUrl = apiUrl;
    this.modelName = modelName;

    this.client = axios.create({
      headers: {
        Authorization: `Bearer ${this.apiKey}`,
        'Content-Type': 'application/json',
      },
    });
  }

  async embedDocuments(texts: Array<string>) {
    if (!texts || texts.length === 0) {
      throw new Error('没有提供文本进行嵌入');
    }

    try {
      console.log(`嵌入 ${texts.length} 个文本`);

      // 批次处理，每批最多处理10个文本
      const batchSize = 10;
      const batches: string[][] = [];

      for (let i = 0; i < texts.length; i += batchSize) {
        batches.push(texts.slice(i, i + batchSize));
      }

      console.log(`将文本分为 ${batches.length} 批进行处理`);

      // 处理每一批
      const embeddings: number[][] = [];
      for (const batch of batches) {
        const batchEmbeddings = await this._embedBatch(batch);
        embeddings.push(...batchEmbeddings);
      }

      return embeddings;
    } catch (error) {
      console.error('嵌入文档失败:', error);
      throw error;
    }
  }

  async embedQuery(text: string) {
    if (!text || text.trim().length === 0) {
      throw new Error('没有提供查询文本进行嵌入');
    }

    try {
      const embeddings = await this._embedBatch([text]);
      return embeddings[0];
    } catch (error) {
      console.error('嵌入查询失败:', error);
      throw error;
    }
  }

  async _embedBatch(batch: string[]): Promise<number[][]> {
    try {
      // 按照硅基流动API的请求格式
      const requestBody = {
        model: this.modelName,
        input: batch,
        encoding_format: 'float',
      };

      console.log(`请求体: ${JSON.stringify(requestBody).substring(0, 100)}...`);

      const response = await this.client.post(this.apiUrl, requestBody);

      if (response.status !== 200) {
        throw new Error(`API请求失败: ${response.status} ${response.statusText}`);
      }

      // 处理硅基流动API的响应格式
      if (
        !response.data ||
        !response.data.data ||
        !response.data.data[0] ||
        !response.data.data[0].embedding
      ) {
        console.error('无效的API响应:', JSON.stringify(response.data).substring(0, 200));
        throw new Error('API响应中没有找到嵌入向量');
      }

      // 返回嵌入向量
      return response.data.data.map((item) => item.embedding);
    } catch (error) {
      console.error('批量嵌入失败:', error);

      // 如果API超时或模型正在加载，等待一段时间后重试
      if (error.response && error.response.status === 503) {
        console.log('模型正在加载，等待5秒后重试...');
        await new Promise((resolve) => setTimeout(resolve, 5000));
        return this._embedBatch(batch);
      }

      throw error;
    }
  }
}
